import copy
import datetime
import os
import json
import torch


class ExperimentLogger:
    """Logs experiment metrics and saves results."""

    def __init__(self, args, run_file_name=None):
        self.start_time = datetime.datetime.now()
        run_file_base = os.path.splitext(run_file_name)[0] if run_file_name else "experiment"

        self.metadata = {
            'dataset': args.dataset,
            'split': 'iid' if args.iid else 'non-iid',
            'num_users': args.num_users,
            'epochs': args.epochs,
            'start_time': self.start_time.strftime("%Y-%m-%d_%H-%M-%S"),
            'run_file': run_file_base
        }
        self.round_data = []
        self.summary = {
            'final_global_accuracy': None,
            'final_local_accuracies': [],
            'model_path': None
        }
        os.makedirs('save', exist_ok=True)

    def log_round(self, round_idx, shapley_vals, weights, duration, local_accs, global_acc):
        """Log metrics for a single training round."""

        def _convert_tensor(value):
            return value.item() if isinstance(value, torch.Tensor) else value

        local_accs = [_convert_tensor(acc) for acc in local_accs]
        global_acc = _convert_tensor(global_acc)

        self.round_data.append({
            'round': round_idx,
            'shapley_values': shapley_vals.tolist(),
            'aggregation_weights': weights.tolist(),
            'time_ms': round(duration * 1000, 2),
            'local_accuracies': [round(acc, 4) for acc in local_accs],
            'global_accuracy': round(global_acc, 4)
        })

    def finalize(self, net_glob, final_local_accs):
        """Save final experiment results and model."""
        final_local_accs = [acc.item() if isinstance(acc, torch.Tensor) else acc
                            for acc in final_local_accs]

        self.summary.update({
            'final_global_accuracy': self.round_data[-1]['global_accuracy'],
            'final_local_accuracies': [round(acc, 4) for acc in final_local_accs]
        })

        base_name = (f"{self.metadata['dataset']}_{self.metadata['split']}_"
                     f"{self.metadata['run_file']}_{self.start_time.strftime('%Y%m%d_%H%M%S')}")

        model_path = os.path.join('save', f"{base_name}_model.pth")
        torch.save(net_glob.state_dict(), model_path)
        self.summary['model_path'] = model_path

        with open(os.path.join('save', f"{base_name}.json"), 'w') as f:
            json.dump({
                'metadata': self.metadata,
                'rounds': self.round_data,
                'summary': self.summary
            }, f, indent=2)

        print(f"\nExperiment data saved to: save/{base_name}.json")
        print(f"Model saved to: {model_path}")